# @Version        : 1.0
# @Update Time    : 2025/3/24 22:36
# @File           : link_track_middleware.py
# @IDE            : PyCharm
# @Desc           : 链路追踪中间件
import uuid
from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.types import ASGIApp

from applications.core.context_manager.track_utils import TrackContextUtils


class LinkTrackMiddleware(BaseHTTPMiddleware):
    def __init__(self, app: ASGIApp, header_name: str = "X-Request-ID"):
        super().__init__(app)
        self.header_name = header_name

    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint):
        incoming = request.headers.get(self.header_name)
        rid = incoming or uuid.uuid4().hex
        TrackContextUtils.set_request_id(rid)
        try:
            response = await call_next(request)
            response.headers[self.header_name] = rid
            return response
        finally:
            # 防止上下文泄漏到复用的 worker 中
            TrackContextUtils.clear_request_id()
